import socket
import threading
import logging
import time
import json
import socket
import ipaddress
from typing import Optional, Callable, Dict, Any, List

# 配置日志
logging.basicConfig(level=logging.INFO,
                    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger('UDPDiscovery')


class UDPDiscoveryBase:
    def __init__(self,
                 port: int = 40000,
                 device_name: str = "UDP Device",
                 discovery_callback: Optional[Callable] = None):
        """
        UDP发现基类 - 提供UDP通信的通用功能

        :param port: UDP通信端口
        :param device_name: 本设备名称
        :param discovery_callback: 发现设备回调函数
        """
        self.port = port
        self.device_name = device_name
        self.discovery_callback = discovery_callback

        self.is_running = False
        self.listen_thread = None
        self.socket = None

        # 存储已发现的设备
        self.discovered_devices = {}
        self.devices_lock = threading.Lock()

        # 子类名称，用于日志和标识
        self.type_name = "Base"

    def _create_socket(self):
        """创建并配置UDP套接字"""
        try:
            # 创建UDP套接字
            sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            sock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
            sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            return sock
        except Exception as e:
            logger.error(f"[{self.type_name}] 创建套接字失败: {str(e)}")
            return None

    def _cleanup_resources(self):
        """清理资源"""
        if self.socket:
            self.socket.close()
            self.socket = None

        with self.devices_lock:
            self.discovered_devices.clear()

    def _safe_json_parse(self, data: bytes) -> Optional[Dict]:
        """安全地解析JSON数据"""
        try:
            return json.loads(data.decode('utf-8', errors='replace'))
        except UnicodeDecodeError:
            logger.warning(f"[{self.type_name}] 接收到无效的UTF-8数据")
            return None
        except json.JSONDecodeError:
            logger.warning(f"[{self.type_name}] 接收到无效的JSON数据")
            return None

    def _safe_json_serialize(self, message: Dict) -> Optional[bytes]:
        """安全地序列化JSON数据"""
        try:
            return json.dumps(message, ensure_ascii=False).encode('utf-8')
        except Exception as e:
            logger.error(f"[{self.type_name}] 序列化JSON数据时出错: {str(e)}")
            return None


class UDPDiscoveryServer(UDPDiscoveryBase):
    def __init__(self,
                 port: int = 40000,
                 device_name: str = "UDP Server",
                 on_client_discovered: Optional[Callable] = None):
        """
        UDP发现服务器类 - 用于监听广播包并响应

        :param port: UDP监听端口
        :param device_name: 本设备名称
        :param on_client_discovered: 发现客户端回调函数，格式: func(client_info: Dict, from_address: tuple)
        """
        # 调用基类构造函数
        super().__init__(port, device_name, on_client_discovered)

        # 重命名基类变量以保持向后兼容性
        self.discovered_clients = self.discovered_devices
        self.clients_lock = self.devices_lock
        self.on_client_discovered = self.discovery_callback

        # 设置类型名称
        self.type_name = "Server"

    def start(self):
        """启动UDP发现服务器"""
        if self.is_running:
            logger.warning(f"[{self.type_name}] UDP发现服务器已经在运行中")
            return

        try:
            # 创建UDP套接字
            self.socket = self._create_socket()
            if not self.socket:
                return

            # 绑定端口，用于接收广播和响应
            self.socket.bind(('0.0.0.0', self.port))

            self.is_running = True

            # 启动监听线程
            self.listen_thread = threading.Thread(target=self._listen, daemon=True)
            self.listen_thread.start()

            logger.info(f"[{self.type_name}] UDP发现服务器已启动，监听端口: {self.port}")

        except Exception as e:
            logger.error(f"[{self.type_name}] 启动UDP发现服务器失败: {str(e)}")
            self.is_running = False
            if self.socket:
                self.socket.close()
                self.socket = None

    def stop(self):
        """停止UDP发现服务器"""
        self.is_running = False

        if self.listen_thread and self.listen_thread.is_alive():
            self.listen_thread.join(timeout=1.0)

        self._cleanup_resources()

        logger.info(f"[{self.type_name}] UDP发现服务器已停止")

    def get_discovered_clients(self) -> Dict[str, Dict[str, Any]]:
        """获取已发现的客户端列表"""
        with self.clients_lock:
            return self.discovered_clients.copy()

    def get_discovered_devices(self) -> Dict[str, Dict[str, Any]]:
        """获取已发现的设备列表（基类方法，保持兼容性）"""
        return self.get_discovered_clients()

    def _listen(self):
        """监听UDP广播消息，优化版本"""
        # 缓冲区大小增加以支持更大的消息
        buffer_size = 4096

        while self.is_running:
            try:
                # 设置超时，让线程可以检查is_running标志
                self.socket.settimeout(1.0)
                data, address = self.socket.recvfrom(buffer_size)

                # 记录接收到的消息源（可选择关闭详细日志）
                logger.debug(f"[{self.type_name}] 接收到来自{address}的消息，长度: {len(data)}字节")

                # 安全的JSON解析
                message = self._safe_json_parse(data)
                if message is None:
                    logger.warning(f"[{self.type_name}] 接收到无效数据来自{address}")
                    continue

                # 基本消息验证
                if not isinstance(message, dict):
                    logger.warning(f"接收到非字典类型的消息来自{address}")
                    continue

                # 获取消息类型
                message_type = message.get('type', 'unknown')

                # 处理发现包
                if message_type == 'discovery':
                    self._process_discovery_message(message, address)

                # 预留其他消息类型的处理
                elif message_type == 'heartbeat':
                    # 未来可以实现心跳消息处理
                    pass

                else:
                    logger.debug(f"接收到未知类型的消息: {message_type}")

            except socket.timeout:
                # 超时是正常的，继续循环
                continue

            except Exception as e:
                if self.is_running:  # 避免在停止服务时打印错误
                    logger.error(f"监听UDP消息时出错: {str(e)}")
                break

    def _process_discovery_message(self, message: dict, address: tuple):
        """处理发现消息的详细逻辑"""
        try:
            # 提取设备信息
            client_info = message.get('device_info', {})
            if not isinstance(client_info, dict):
                logger.warning(f"[{self.type_name}] 无效的device_info格式来自{address}")
                client_info = {}

            # 优先使用客户端提供的client_ip（如果存在且有效）
            # 这解决了NAT环境或多网卡环境下的IP识别问题
            reported_client_ip = client_info.get('client_ip')
            source_ip = address[0]

            # 使用客户端报告的IP，但记录来源IP以便参考
            client_ip = reported_client_ip if reported_client_ip else source_ip

            # 获取设备名称和其他信息
            device_name = client_info.get('name', 'Unknown Device')
            timestamp = client_info.get('timestamp', time.time())
            c = client_info.get("client_ip", "")
            # 构建完整的客户端信息
            updated_client_info = {
                'ip': client_ip,
                'source_ip': source_ip,  # 记录原始来源IP
                'name': device_name,
                'timestamp': time.time(),
                'reported_timestamp': timestamp,
                'last_seen': time.time()
            }

            # 如果客户端提供了其他信息，也一并保存
            for key, value in client_info.items():
                if key not in ['name', 'timestamp', 'client_ip']:
                    updated_client_info[key] = value

            # 安全地更新已发现客户端列表
            with self.clients_lock:
                client_key = f"{client_ip}:{self.port}"
                self.discovered_clients[client_key] = updated_client_info

                # 清理过期的客户端（超过30秒未更新）
                self._cleanup_expired_clients()

            # 记录发现新客户端
            logger.debug(f"[{self.type_name}] 发现客户端: {device_name} ({client_ip})，来源IP: {source_ip}")

            # 调用回调函数（如果有）
            if self.on_client_discovered:
                try:
                    # 传递完整的客户端信息和原始地址
                    self.on_client_discovered(updated_client_info, address)
                except Exception as e:
                    logger.error(f"[{self.type_name}] 回调函数执行出错: {str(e)}")

            # 发送响应包
            self._send_response(address)

        except Exception as e:
            logger.error(f"[{self.type_name}] 处理发现消息时出错: {str(e)}")

    def _cleanup_expired_clients(self):
        """清理过期的客户端信息"""
        current_time = time.time()
        expiration_time = 30  # 30秒过期

        # 使用列表推导式创建要保留的客户端列表
        valid_clients = {}
        for key, client in self.discovered_clients.items():
            if current_time - client.get('last_seen', 0) < expiration_time:
                valid_clients[key] = client

        # 只在有变化时更新字典
        if len(valid_clients) != len(self.discovered_clients):
            self.discovered_clients = valid_clients
            logger.debug(f"[{self.type_name}] 清理过期客户端，剩余: {len(valid_clients)}个客户端")

    def _send_response(self, target_address: tuple):
        """发送响应消息给指定地址，优化版本"""
        if not self.is_running or not self.socket:
            logger.warning(f"[{self.type_name}] 尝试发送响应，但服务未运行或套接字未初始化")
            return

        try:
            # 验证目标地址
            if not isinstance(target_address, tuple) or len(target_address) != 2:
                logger.warning(f"[{self.type_name}] 无效的目标地址: {target_address}")
                return

            # 构建更完整的响应消息
            response_message = {
                'type': 'response',
                'version': '1.0',  # 添加版本信息，便于未来扩展
                'device_info': {
                    'name': self.device_name,
                    'timestamp': time.time(),
                    'server_ip': f"{socket.gethostname()}",
                    'server_port': f"{self.port}"
                }
            }

            # 安全的JSON序列化
            message_data = self._safe_json_serialize(response_message)
            if message_data is None:
                return

            # 检查消息大小，避免过大的数据包
            if len(message_data) > 8192:  # 限制消息大小为8KB
                logger.warning(f"响应消息过大 ({len(message_data)}字节)，已忽略")
                return

            # 发送响应消息
            bytes_sent = self.socket.sendto(message_data, target_address)
            logger.debug(f"向{target_address}发送响应消息，大小: {bytes_sent}字节")

        except socket.error as sock_error:
            # 更详细的套接字错误处理
            error_code = getattr(sock_error, 'errno', None)
            if error_code == 10054:  # Connection reset by peer
                logger.debug(f"发送响应到{target_address}时连接被重置")
            elif error_code == 10065:  # No route to host
                logger.warning(f"无法路由到主机: {target_address}")
            else:
                logger.error(f"发送响应消息到{target_address}时出错: {str(sock_error)}")
        except Exception as e:
            logger.error(f"发送响应消息时发生未预期错误: {str(e)}")


class UDPDiscoveryClient(UDPDiscoveryBase):
    def __init__(self,
                 port: int = 40000,
                 broadcast_interval: int = 5,
                 device_name: str = "UDP Client",
                 on_server_discovered: Optional[Callable] = None):
        """
        UDP发现客户端类 - 用于发送广播包并接收响应

        :param port: UDP通信端口
        :param broadcast_interval: 广播间隔(秒)
        :param device_name: 本设备名称
        :param on_server_discovered: 发现服务器回调函数，格式: func(server_info: Dict, from_address: tuple)
        """
        # 调用基类构造函数
        super().__init__(port, device_name, on_server_discovered)

        # 子类特有的参数
        self.broadcast_interval = broadcast_interval
        self.broadcast_thread = None

        # 重命名基类变量以保持向后兼容性
        self.discovered_servers = self.discovered_devices
        self.servers_lock = self.devices_lock
        self.on_server_discovered = self.discovery_callback

        # 设置类型名称
        self.type_name = "Client"

        # 尝试导入netifaces库，如果没有则使用备选方案
        self.use_netifaces = False
        try:
            import netifaces
            self.netifaces = netifaces
            self.use_netifaces = True
        except ImportError:
            logger.warning(f"[{self.type_name}] 未安装netifaces库，将使用简化的广播地址获取方法")

    def start(self):
        """启动UDP发现客户端"""
        if self.is_running:
            logger.warning(f"[{self.type_name}] UDP发现客户端已经在运行中")
            return

        try:
            # 创建UDP套接字
            self.socket = self._create_socket()
            if not self.socket:
                return

            self.is_running = True

            # 启动监听线程，用于接收响应
            self.listen_thread = threading.Thread(target=self._listen, daemon=True)
            self.listen_thread.start()

            logger.info(f"[{self.type_name}] UDP发现客户端已启动")

        except Exception as e:
            logger.error(f"[{self.type_name}] 启动UDP发现客户端失败: {str(e)}")
            self.is_running = False
            if self.socket:
                self.socket.close()
                self.socket = None

    def stop(self):
        """停止UDP发现客户端"""
        self.is_running = False

        if self.broadcast_thread and self.broadcast_thread.is_alive():
            self.broadcast_thread.join(timeout=1.0)

        if self.listen_thread and self.listen_thread.is_alive():
            self.listen_thread.join(timeout=1.0)

        self._cleanup_resources()

        logger.info(f"[{self.type_name}] UDP发现客户端已停止")

    def start_broadcasting(self):
        """开始广播设备信息"""
        if not self.is_running:
            logger.warning("UDP发现客户端未启动，请先调用start()方法")
            return

        if self.broadcast_thread and self.broadcast_thread.is_alive():
            logger.warning("广播线程已经在运行中")
            return

        self.broadcast_thread = threading.Thread(target=self._broadcast, daemon=True)
        self.broadcast_thread.start()
        logger.info("开始广播设备信息")

    def broadcasting_once(self):
        self._broadcast_function()

    def stop_broadcasting(self):
        """停止广播设备信息"""
        # 广播线程会检查is_running标志，所以这里不需要额外操作
        logger.info("停止广播设备信息")

    def get_discovered_servers(self) -> Dict[str, Dict[str, Any]]:
        """获取已发现的服务器列表"""
        with self.servers_lock:
            return self.discovered_servers.copy()

    def get_discovered_devices(self) -> Dict[str, Dict[str, Any]]:
        """获取已发现的设备列表（基类方法，保持兼容性）"""
        return self.get_discovered_servers()

    def _listen(self):
        """监听UDP响应消息"""
        while self.is_running:
            try:
                # 设置超时，让线程可以检查is_running标志
                self.socket.settimeout(1.0)
                data, address = self.socket.recvfrom(1024)

                # 安全的JSON解析
                message = self._safe_json_parse(data)
                if message is None:
                    logger.warning(f"[{self.type_name}] 接收到无效数据来自{address}")
                    continue

                # 处理响应包
                if message.get('type') == 'response':
                    server_info = message.get('device_info', {})
                    server_ip = address[0]

                    # 更新已发现服务器列表
                    with self.servers_lock:
                        server_key = f"{server_ip}:{self.port}"
                        self.discovered_servers[server_key] = {
                            'ip': server_ip,
                            'name': server_info.get('name', 'Unknown'),
                            'server_ip': server_info.get('server_ip', 'Unknown'),
                            'server_port': server_info.get('server_port', 'Unknown'),
                            'timestamp': time.time()
                        }

                    # 调用回调函数
                    if self.on_server_discovered:
                        try:
                            self.on_server_discovered(
                                self.discovered_servers[server_key],
                                address
                            )
                        except Exception as e:
                            logger.error(f"[{self.type_name}] 回调函数执行出错: {str(e)}")

            except socket.timeout:
                # 超时是正常的，继续循环
                continue

            except Exception as e:
                if self.is_running:  # 避免在停止服务时打印错误
                    logger.error(f"[{self.type_name}] 监听UDP消息时出错: {str(e)}")
                break

    def _get_broadcast_addresses(self) -> List[str]:
        """获取所有网络接口的广播地址"""
        broadcast_addresses = []

        try:
            if self.use_netifaces:
                # 使用netifaces库获取所有网络接口和广播地址
                for iface in self.netifaces.interfaces():
                    try:
                        # 获取接口的IPv4地址信息
                        addrs = self.netifaces.ifaddresses(iface)
                        if self.netifaces.AF_INET in addrs:
                            for addr_info in addrs[self.netifaces.AF_INET]:
                                if 'broadcast' in addr_info:
                                    broadcast_addresses.append(addr_info['broadcast'])
                    except Exception as e:
                        logger.debug(f"[{self.type_name}] 处理网络接口{iface}时出错: {str(e)}")
            else:
                # 备选方案：获取主机名对应的IP地址，并构造广播地址
                # 获取主机名
                hostname = socket.gethostname()
                # 获取主机的所有IP地址
                addr_info = socket.getaddrinfo(hostname, None, socket.AF_INET)

                # 提取IP地址并构造广播地址
                for info in addr_info:
                    ip = info[4][0]
                    try:
                        # 假设子网掩码为255.255.255.0构造广播地址
                        # 这是一个简化的方法，可能不适用于所有网络配置
                        ip_parts = ip.split('.')
                        if len(ip_parts) == 4:
                            broadcast_ip = f"{ip_parts[0]}.{ip_parts[1]}.{ip_parts[2]}.255"
                            broadcast_addresses.append(broadcast_ip)
                    except Exception as e:
                        logger.debug(f"[{self.type_name}] 构造广播地址时出错: {str(e)}")

                # 添加默认的广播地址
                if not broadcast_addresses:
                    broadcast_addresses.append('<broadcast>')

            # 去重并确保至少有一个广播地址
            broadcast_addresses = list(set(broadcast_addresses))
            logger.debug(f"[{self.type_name}] 发现的广播地址: {broadcast_addresses}")
        except Exception as e:
            logger.error(f"[{self.type_name}] 获取广播地址时出错: {str(e)}")
            # 出错时使用默认广播地址
            broadcast_addresses = ['<broadcast>']

        return broadcast_addresses

    def _get_client_ip(self, broadcast_address: str) -> str:
        """根据广播地址获取对应的网卡IP地址

        :param broadcast_address: 广播地址
        :return: 对应的网卡IP地址
        """
        try:
            if self.use_netifaces:
                # 使用netifaces库精确匹配广播地址对应的网卡IP
                for iface in self.netifaces.interfaces():
                    try:
                        addrs = self.netifaces.ifaddresses(iface)
                        if self.netifaces.AF_INET in addrs:
                            for addr_info in addrs[self.netifaces.AF_INET]:
                                if 'broadcast' in addr_info and addr_info['broadcast'] == broadcast_address:
                                    return addr_info['addr']
                    except Exception as e:
                        logger.debug(f"[{self.type_name}] 处理网络接口{iface}时出错: {str(e)}")

            # 如果使用netifaces库没有找到匹配的IP，或者没有安装netifaces库
            # 使用简化的方法：从广播地址推断网卡IP
            if broadcast_address != '<broadcast>':
                try:
                    # 假设子网掩码为255.255.255.0
                    # 例如：如果广播地址是192.168.1.255，那么网卡IP可能是192.168.1.x
                    broadcast_parts = broadcast_address.split('.')
                    if len(broadcast_parts) == 4:
                        # 尝试获取主机的所有IP地址，然后找出属于同一网段的
                        hostname = socket.gethostname()
                        addr_info = socket.getaddrinfo(hostname, None, socket.AF_INET)

                        for info in addr_info:
                            ip = info[4][0]
                            ip_parts = ip.split('.')
                            if len(ip_parts) == 4 and ip_parts[0] == broadcast_parts[0] and ip_parts[1] == \
                                    broadcast_parts[1] and ip_parts[2] == broadcast_parts[2]:
                                return ip
                except Exception as e:
                    logger.debug(f"[{self.type_name}] 根据广播地址推断网卡IP时出错: {str(e)}")

            # 如果以上方法都失败，返回默认的本地IP
            return '127.0.0.1'
        except Exception as e:
            logger.error(f"[{self.type_name}] 获取客户端IP时出错: {str(e)}")
            return '127.0.0.1'

    def _build_broadcast_ip_map(self, broadcast_addresses: List[str]) -> Dict[str, str]:
        """构建广播地址到网卡IP的映射字典

        :param broadcast_addresses: 广播地址列表
        :return: 广播地址到网卡IP的映射字典
        """
        broadcast_ip_map = {}
        try:
            for broadcast_address in broadcast_addresses:
                # 为每个广播地址获取对应的网卡IP
                client_ip = self._get_client_ip(broadcast_address)
                broadcast_ip_map[broadcast_address] = client_ip
                logger.debug(f"[{self.type_name}] 映射广播地址{broadcast_address}到网卡IP: {client_ip}")
        except Exception as e:
            logger.error(f"[{self.type_name}] 构建广播地址到网卡IP的映射时出错: {str(e)}")

        return broadcast_ip_map

    def _broadcast(self):
        """广播设备发现消息"""
        while self.is_running:
            self._broadcast_function()

    def _broadcast_function(self):

        try:
            # 获取所有网络接口的广播地址
            broadcast_addresses = self._get_broadcast_addresses()

            # 在循环外预先构建广播地址到网卡IP的映射字典
            broadcast_ip_map = self._build_broadcast_ip_map(broadcast_addresses)

            # 当前时间戳，避免在循环中重复生成
            current_timestamp = time.time()

            # 对每个广播地址发送广播消息
            for broadcast_address in broadcast_addresses:
                try:
                    # 从映射字典中获取对应的网卡IP
                    client_ip = broadcast_ip_map.get(broadcast_address, '127.0.0.1')

                    # 构建发现消息
                    discovery_message = {
                        'type': 'discovery',
                        'device_info': {
                            'name': self.device_name,
                            'timestamp': current_timestamp,
                            'client_ip': client_ip
                        }
                    }
                    message_data = self._safe_json_serialize(discovery_message)
                    if message_data is None:
                        continue

                    self.socket.sendto(message_data, (broadcast_address, self.port))
                    logger.debug(
                        f"[{self.type_name}] 向{broadcast_address}:{self.port}发送广播消息，使用网卡IP: {client_ip}")
                except Exception as e:
                    logger.warning(
                        f"[{self.type_name}] 向{broadcast_address}:{self.port}发送广播消息时出错: {str(e)}")

            # 等待指定的广播间隔
            for _ in range(int(self.broadcast_interval * 10)):
                if not self.is_running:
                    break
                time.sleep(0.1)

        except Exception as e:
            if self.is_running:  # 避免在停止服务时打印错误
                logger.error(f"广播过程中出错: {str(e)}")
            time.sleep(1.0)  # 出错时暂停一秒再试


# 示例用法
if __name__ == '__main__':
    # 服务器端示例
    def on_client_found(client_info, from_address):
        logger.info(f"服务器发现客户端: {client_info['name']} (IP: {client_info['ip']}) 来自 {from_address}")
        # logger.info(json.dumps(client_info, ensure_ascii=False, indent=2))


    try:
        # 创建并启动UDP发现服务器
        server = UDPDiscoveryServer(
            port=40000,
            device_name="UDP Server",
            on_client_discovered=on_client_found
        )
        server.start()

        logger.info("UDP发现服务已启动，按Ctrl+C停止")

        # 主循环，定期打印已发现的设备
        while True:
            # # 打印服务器发现的客户端
            # clients = server.get_discovered_clients()
            # if clients:
            #     print(f"\n服务器已发现的客户端列表 ({len(clients)}):")
            #     for client_key, client_info in clients.items():
            #         print(f"- {client_info['name']} (IP: {client_info['ip']})")
            #
            # # 打印客户端发现的服务器
            # servers = client.get_discovered_servers()
            # if servers:
            #     print(f"\n客户端已发现的服务器列表 ({len(servers)}):")
            #     for server_key, server_info in servers.items():
            #         print(f"- {server_info['name']} (IP: {server_info['ip']})")
            #
            time.sleep(5)

    except KeyboardInterrupt:
        print("正在停止UDP发现服务...")

    finally:
        # 停止服务
        if 'server' in locals():
            server.stop()
        print("UDP发现服务已停止")
